# 此函数用于对称分布傅里叶的色散部分的仿真，GVD+PMD
# 传入参数：
#   polar_n:偏振数
#   beta_2:群速度色散,s^2/km
#   f_omega
#   deltaz:每个span进行分布傅里叶的长度
#   PMD:是否添加偏振模色散
#   PMD_str_para:PDM构造好的参数

import numpy as np
from scipy.io import loadmat
from numpy.fft import fft, ifft ,fftshift, ifftshift

def Dispersion(signal, Signal, Fiber_para, PMD_para, Dis):
    data = signal
    data_fft = fftshift(fft(data))
    #1/2衰减后的时域信号

    '=================================================信号相关参数'
    npol = Signal['polarization']
    lamda_wave = Signal['lamda_wave']
    field = Signal['field']
    C = Signal['C']
    f_omega = Signal['f_omega']

    '================================================光纤相关参数'
    L = Fiber_para['L']
    lamda_Fiber = Fiber_para['lambda']
    StepSize = Fiber_para['step_size']

    '================================================差分群时延相关参数'
    # 群速度及高阶色散参数
    beta_2_fiber = Fiber_para['beta_2']  # [s^2/km]
    beta_3_fiber = Fiber_para['beta_3']  # [s^3/km]

    # 主体部分
    if npol == 1:

        # single-polarization
        beta_2 = beta_2_fiber  # beta2 [s^2/km]
        beta_3 = beta_3_fiber  # [s^3/km]
        beta_ = f_omega * (f_omega * (beta_2 / 2 + f_omega * beta_3 / 6))

        # 色散[1 / km]
        data_fft = data_fft * np.exp(-1j * beta_ * StepSize / 2)
        sig = ifft(ifftshift(data_fft))

    return sig